package com.learn.grpc.loadbalance;

import io.grpc.ConnectivityState;
import io.grpc.ConnectivityStateInfo;
import io.grpc.LoadBalancer;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.stream.Collectors;

import static com.learn.grpc.loadbalance.LbLoadBalancer.STATE_INFO;
import static io.grpc.ConnectivityState.*;

/** 监听 channel 状态
 *  channel 状态变更时，同步更新负载均衡器对应的 实例列表
 */
@Slf4j
public class LbSubChannelStateListener implements LoadBalancer.SubchannelStateListener {
    private final LoadBalancer.Subchannel subChannel;
    private final LoadBalancer.Helper helper;
    private final LbLoadBalancer loadBalancer;

    public LbSubChannelStateListener(LbLoadBalancer lbLoadBalancer,
                                     LoadBalancer.Subchannel subChannel,
                                     LoadBalancer.Helper helper) {
        this.loadBalancer = lbLoadBalancer;
        this.subChannel = subChannel;
        this.helper = helper;
    }

    @Override
    public void onSubchannelState(ConnectivityStateInfo stateInfo) {

        // subChannel 连接状态
        LbRef<ConnectivityState> subChannelConnectState = subChannel.getAttributes().get(STATE_INFO);

        ConnectivityState currentState = subChannelConnectState.getValue();
        ConnectivityState newState = stateInfo.getState();

        // 如果 channel 已关闭
        if (newState == SHUTDOWN) {
            InetSocketAddress socketAddress = (InetSocketAddress) subChannel.getAddresses().getAddresses().get(0);
            log.info("关闭channel:{}:{}" , socketAddress.getAddress(), socketAddress.getPort());
            return;
        }

        // 如果当前状态为
        if (currentState == TRANSIENT_FAILURE) {
            if (newState == CONNECTING || newState == IDLE) {
                log.error("建立连接或者失败 {}" , subChannel.getAddresses().getAddresses());
                return;
            }
        }

        //主要逻辑: 如果 channel 就绪，则尝试建立连接
        if (newState == READY) {
            subChannel.requestConnection();
        }

        // 修改当前连接状态
        subChannelConnectState.setValue(newState);

        // 有任意 channel 状态变更，都对负载均衡器进行更新
        updateLoadBalancer();
    }

    // 更新负责均衡器
    private void updateLoadBalancer() {
        // 筛选可用 channel 列表, 可用指的是 ready 状态的 channel
        List<LoadBalancer.Subchannel> readySubChannels = loadBalancer.getSubChannelMap()
             .values()
             .stream()
             .filter(s -> s.getAttributes().get(STATE_INFO).getValue() == READY)
             .collect(Collectors.toList());

        // 注册 channel 选择器
        if (readySubChannels.isEmpty()) {
            helper.updateBalancingState(CONNECTING, new LbSubChannelPicker(null));
        } else {
            helper.updateBalancingState(READY, new LbSubChannelPicker(readySubChannels));
        }
    }
}

